#include <WProgram.h>
#include "ADB.h"

#define DEBUG
#define TRACE
#include "debug.h"


ADBconnection::ADBconnection(char* destination) 
  : usb(),
    usbConnected(false), status(CLOSED), 
    localID(1), remoteID(0),
    msgCmd(0),msgRemoteId(0),msgDataLen(0),msgDataCRC(0)
{
  strncpy(ADBconnection::destination, destination, 
	  sizeof(ADBconnection::destination)-1);
}

/*
  init()

  Return true if USB connection to ADB device made.
  read descriptors and parse.

 */
bool ADBconnection::init()
{
  if (status == HOSTCONNECT) {
    return true;
  }

  _DTR();
  initUsb();
  if (!usbConfiguring()) {
    _DPRINT("USB config.");
    return false;
  }

  delay(100);
  uint8_t numConfig = numDeviceDescriptors();
  if (numConfig == 0) {
    _DPRINT("No device descr.");
    return false;
  }

  _DTR();
  // interested only in first config.
  if (!foundADB()) {
    _DPRINT("No ADB device.");
    return false;
  }

  if (!setDevConfig()) {
    _DPRINT("Host config.");
    return false;
  }
  usbConnected = true;

  _DTR();
  if (!hostConnect()) {
    _DPRINT("Host connection.");
    return false;
  }
  status = HOSTCONNECT;

  return true;
}

/*
   init() helpers
 */
void ADBconnection::initUsb()
{
  usb.handle.powerOn();
}

bool ADBconnection::usbConfiguring()
{
  uint32_t deadline = millis() + ADB_HOST_CONNECTION_WAIT_TIME;
  while (millis() < deadline) {
    usb.handle.MAX3421E::Task();
    usb.handle.USB::Task();
    if (usb.handle.getUsbTaskState() >= USB_STATE_CONFIGURING) {
      return true;
    }
  }
  _DPRINT3("state",usb.handle.getUsbTaskState(),HEX);
  return false;
}

uint8_t ADBconnection::numDeviceDescriptors()
{
  USB_DEVICE_DESCRIPTOR buf;
  byte rcode = usb.handle.getDevDescr(1, 0, sizeof(buf), (char*)&buf);
  if (rcode) {
    _DPRINT3("getDevDescr",rcode,HEX);
    return 0;
  }
  return buf.bNumConfigurations;
}

// Only look at the first config, i.e. config = 0
bool ADBconnection::foundADB()
{
  char* bufp = dataBuf;
  byte rcode;
  byte descLen;
  byte descType;
  uint16_t totalLen;

  rcode = usb.handle.getConfDescr(1, 0, 4, 0, dataBuf);  //get total length
  if (rcode) {
    _DPRINT3("getConfDescr",rcode,HEX);
    return false;
  }

  totalLen = dataBuf[2] + (dataBuf[3]<<8);
  if( totalLen > sizeof(dataBuf) ) {
    totalLen = sizeof(dataBuf);
  }

  rcode = usb.handle.getConfDescr(1, 0, totalLen, 0, dataBuf ); //get the descriptor

  // descriptors are a configuration descriptor listing the number of interfaces 
  // followed by each interface. Each interface has an interface descriptor followed
  // by 2 endpoint descriptors. We can skip data if the interface descriptor doesn't
  // indicate an ADB interface.

  for (; bufp < dataBuf+totalLen; bufp += descLen) { 
    descLen = *bufp;
    descType = *(bufp+1);
    switch( descType ) {
      case USB_DESCRIPTOR_CONFIGURATION:
	getDescrConfig((USB_CONFIGURATION_DESCRIPTOR*)bufp);
        break;
      case USB_DESCRIPTOR_INTERFACE:
	if (isADBInterface((USB_INTERFACE_DESCRIPTOR*)bufp)) {
	  return getADBDescrData(bufp);
	}
        break;
      default:
        break;
    }
  }

  return false;
}

void ADBconnection::getDescrConfig(USB_CONFIGURATION_DESCRIPTOR* cDesc)
{
  usb.config = cDesc->bConfigurationValue;
}

bool ADBconnection::isADBInterface(USB_INTERFACE_DESCRIPTOR* iDesc)
{
  if (iDesc->bNumEndpoints!=2 ||
      iDesc->bInterfaceProtocol != ADB_PROTOCOL ||
      iDesc->bInterfaceClass != ADB_CLASS || 
      iDesc->bInterfaceSubClass != ADB_SUBCLASS) {
    return false;
  }

  return true;
}

// Get the interface & endpoint data following this ADB interface descriptor
bool ADBconnection::getADBDescrData(char* bufp)
{
  uint8_t descLen = *bufp;
  uint8_t descType = *(bufp+1);

  if (descType != USB_DESCRIPTOR_INTERFACE) {
    _DPRINT3("interface",descType,HEX);
    return false;
  }
  usb.interface = ((USB_INTERFACE_DESCRIPTOR*)bufp)->bInterfaceNumber;
  bufp += descLen;

  for (uint8_t numEPs=1, descLen=*bufp; numEPs<=2; numEPs++,bufp+=descLen) {

    descType = *(bufp+1);
    if (descType != USB_DESCRIPTOR_ENDPOINT) {
      _DPRINT3("endpoint",descType,HEX);
      return false;
    }
  
    USB_ENDPOINT_DESCRIPTOR* epp = (USB_ENDPOINT_DESCRIPTOR*)bufp;
    uint8_t adr = epp->bEndpointAddress;
    if (adr & 0x80) {
      usb.inputEP = (adr & 0x0F);
    } else {
      usb.outputEP = (adr & 0x0F);
    }

  }
  return true;
}


bool ADBconnection::setDevConfig()
{
  if (usb.outputEP >= MAX_USB_EPS ||
      usb.inputEP >= MAX_USB_EPS) {
    Serial.println("MAX_USB_EPS too low. Aborting."); while(true);
  }

  // TODO: decipher this.
  usb.epRec[ 0 ] = *( usb.handle.getDevTableEntry( 0,0 ));
  usb.epRec[1].MaxPktSize = 8;
  usb.epRec[1].sndToggle = bmSNDTOG0;
  usb.epRec[1].rcvToggle = bmRCVTOG0;

  usb.epRec[usb.inputEP].MaxPktSize = 
    usb.epRec[usb.outputEP].MaxPktSize = 8;
  usb.epRec[usb.inputEP].sndToggle = 
    usb.epRec[usb.outputEP].sndToggle = bmSNDTOG0;
  usb.epRec[usb.inputEP].rcvToggle = 
    usb.epRec[usb.outputEP].rcvToggle = bmRCVTOG0;
  usb.handle.setDevTableEntry( 1, usb.epRec );

  // TODO: what endpoint? 0 is control ep?
  uint8_t rcode = usb.handle.setConf(1, 0, usb.config);
  if (rcode) {
    _DPRINT3("Host connect send.",rcode,HEX);
    return false;
  }

  return true;
}


/*
  hostConnect()
  Make host level connection within which software-level connections are
  supported by ADB.
 */
bool ADBconnection::hostConnect()
{
  if (!sendHostConnect()) {
    return false;
  }
  _DTR();

  uint32_t deadline = millis() + ADB_HOST_CONNECTION_WAIT_TIME;
  while (millis() < deadline) {
    usb.handle.MAX3421E::Task();
    usb.handle.USB::Task();
    if (recvMessage() && recvdHostConnectReply()) {
      return true;
    }
  }
  return false;
}

// Set states based on rcode.
bool ADBconnection::handleUsbErr(char* where, uint8_t rcode)
{
  switch(rcode) {
  case hrSUCCESS:
    return true;

  case hrNAK:
    return false;		// nothing to report.

  // failure cases
  case hrJERR:			// fall thru to close connection
  case hrTIMEOUT:
    close();
    status=CLOSED;		// assume HOST connection is also lost.
    _DPRINT("close host conn.");
  default:
    _DPRINT3(where,rcode,HEX);
    return false;
  }
  return false;
}

bool ADBconnection::sendHostConnect()
{
  char * id = "host::arduino";
  return send(A_CNXN,0x1000000,ADB_MAX_PAYLOAD,strlen(id)+1, (uint8_t*)id);
}

bool ADBconnection::send(uint32_t command, uint32_t arg0, uint32_t arg1,
			 uint32_t len, uint8_t* data) 
{
  ADBmessage message;
  uint32_t count, sum = 0;
  uint8_t * x;
  uint8_t rcode;

  // Calculate data checksum
  count = len;
  x = data;
  while(count-- > 0) sum += *x++;

  // Fill out the message record.
  message.command = command;
  message.arg0 = arg0;
  message.arg1 = arg1;
  message.dataLen = len;
  message.dataCRC = sum;
  message.magic = command ^ 0xffffffff;

  rcode = usb.handle.outTransfer(1,usb.outputEP, sizeof(ADBmessage), 
				 (char*)&message);
  if (rcode) return handleUsbErr("send-hdr",rcode);
  if (!data) return true;
  return handleUsbErr("send-data",
		      usb.handle.outTransfer(1, usb.outputEP, len, (char*)data));
}


// Is last message a hostConnect reply? If so, handle it.
bool ADBconnection::recvdHostConnectReply()
{
  if (msgCmd != A_CNXN) {
    return false;
  }
  // don't need data accompanying message but need to get it.
  if (!recvData(sizeof(dataBuf),(uint8_t*)dataBuf)) {
    return false;
  }

  msgCmd = 0;
  msgDataLen = 0;
  return true;
}


// receive data specified by message.
bool ADBconnection::recvData(uint16_t len, uint8_t* buf)
{
  if (msgDataLen > len) {
    _DPRINT3("dataBuf len.",msgDataLen,DEC);
    return false;
  }

  uint16_t dataLen;
  uint8_t rcode;
  if (rcode=usb.handle.inTransferX(1, usb.inputEP, msgDataLen, (char*)buf, &dataLen)) {
    return handleUsbErr("data read",rcode);
  }

  // verify the data CRC
  uint32_t crc = 0;
  for (uint16_t i=0;i<msgDataLen;i++) {
    crc += buf[i];
  }

  if (crc != msgDataCRC) {
    _DPRINT("crc");
    return false;
  }

  // send peer an OKAY message.
  sendOK();

  return true;
}

/*
  recvMessage()
  See if data is there to be read. 
  Returns :
   true if yes.
   false if no

 */
bool ADBconnection::recvMessage()
{
  ADBmessage msg;

  uint16_t dataLen;
  uint8_t rcode = usb.handle.inTransferX(1, usb.inputEP, sizeof(msg), (char*)&msg, &dataLen, 1);
  if (rcode) {
    return handleUsbErr("inData",rcode);
  }

  // If the message is corrupt, return.
  if (dataLen != sizeof(ADBmessage) ||
      msg.magic != (msg.command ^ 0xffffffff)) {
    _DPRINT("Msg corrupt"); 
    _DPRINT3("dataLen",dataLen,HEX);
    _DPRINT3("cmd",msg.command,HEX);
    _DPRINT3("magic",msg.magic,HEX);
    return false;
  }

  msgCmd = msg.command;
  msgRemoteId = msg.arg0;
  msgDataLen = msg.dataLen;
  msgDataCRC = msg.dataCRC;

  _DPRINT3("msgCmd:", msgCmd, HEX);
  _DPRINT3("msgDataLen:",msgDataLen,HEX);
  _DPRINT3("Status:",status,HEX);
  return true;
}

/*
  connect()
  Make the ADB connection & return status.

*/
bool ADBconnection::connect()
{
  if (status >= OPEN) return true;
  
  if (status < HOSTCONNECT) {
    return false;
  }

  if (!send(A_OPEN,localID,remoteID,strlen(destination)+1,
	    (uint8_t*)destination)) {
    return false;
  }
  
  uint32_t deadline = millis() + ADB_CONNECTION_WAIT_TIME;
  while (millis() < deadline) {
    usb.handle.MAX3421E::Task();
    usb.handle.USB::Task();
    if (recvMessage() && recvdConnectReply()) {
      return true;
    }
  }

  return false;
}

// Handle reply for connection request
bool ADBconnection::recvdConnectReply()
{
  if (msgCmd != A_OKAY) {
    return false;
  }

  remoteID = msgRemoteId;
  status = OPEN;
  msgCmd = 0;
  msgDataLen = 0;
  return true;
}

// Write without confirmation. async
bool ADBconnection::write(uint8_t * data, uint16_t length)
{
  if (status < OPEN) {
    _DPRINT("open");
    return false;
  }

  uint8_t rcode;
  if (!send(A_WRTE,localID,remoteID,length,data)) {
    return false;
  }
  status = WRITING;
  return true;
}
  
// Write and wait for confirmation. sync.
bool ADBconnection::writeSync(uint8_t * data, uint16_t length)
{
  if (!write(data, length)) {
    return false;
  }

  uint32_t deadline = millis() + ADB_WRITE_WAIT_TIME;
  while (millis() < deadline) {
    usb.handle.MAX3421E::Task();
    usb.handle.USB::Task();
    if (recvMessage() && recvdWriteOK()) {
      return true;
    }
  }

  return false;
  
}

// Handle reply for write request
bool ADBconnection::recvdWriteOK()
{
  if (msgCmd != A_OKAY) {
    return false;
  }

  status = OPEN;
  msgCmd = 0;
  msgDataLen = 0;
  return true;
}

// call periodically. To let the underlying USB layer do its thing.
// To open the connection if closed & look for writes to us.
void ADBconnection::poll()
{
  usb.handle.MAX3421E::Task();
  usb.handle.USB::Task();
  
  connect();
  checkForMsgs();
  connect();			// for good measure.
}

// Ignore OKAYs. Return if WRTE => host is writing us.
void ADBconnection::checkForMsgs()
{
  while (recvMessage()) {
    switch(msgCmd) {
    case A_CLSE:
      status = HOSTCONNECT;
      break;
    case A_OKAY:
      if (status == WRITING) {
	status = OPEN;
      } else {
	// if we write async, it is possible to get these.
	//_DPRINT3("A_OKAY unexp",status,DEC);
      }
      break;
    case A_WRTE:
      status = RECEIVING;
      return;

    default:
      _DPRINT3("ADB cmd",msgCmd,HEX);
    }
  }
}

uint16_t ADBconnection::available()
{
  /*
  if (msgDataLen) {
    _DPRINT3("msgCmd",msgCmd,HEX);
  }
  */
  if (status != RECEIVING) return 0;
  return msgDataLen;
}

// Demand they read the entire data.
int ADBconnection::read(uint16_t len, uint8_t* buf)
{
  if (status != RECEIVING) {
    _DPRINT3("status",status,DEC);
    return -1;
  }
  if (!msgDataLen || len != msgDataLen) {
    _DPRINT3("msgDataLen",msgDataLen,HEX);
    return 0;
  }

  if (!recvData(len,buf)) return -2;

  status = OPEN;
  msgCmd = 0;
  msgDataLen = 0;
  return len;
}

void ADBconnection::sendOK()
{
  (void)send(A_OKAY,localID,remoteID,0,NULL);
}

void ADBconnection::close()
{
  (void)send(A_CLSE,localID,remoteID,0,NULL);
}
